# coding=utf-8
import math
import numpy as np
import tensorflow as tf
slim = tf.contrib.slim
import a11_net_params as netparam
import a13_gen_anchor_box

def bboxes_decode(feat_localizations, 
                  anchors,
                  scope='ssd_bboxes_decode'):
    return tf_ssd_bboxes_decode(
		feat_localizations, 
		anchors,
		prior_scaling = netparam.default_params.prior_scaling,
		scope = scope)

"""
Compute the relative bounding boxes from the SSD net features and
reference anchors bounding boxes.

Arguments:
  feat_localizations: List of Tensors containing localization features.
  anchors: List of numpy array containing anchor boxes.

Return:
  List of Tensors Nx4: ymin, xmin, ymax, xmax
"""
def tf_ssd_bboxes_decode(feat_localizations,
                         anchors,
                         prior_scaling=[0.1, 0.1, 0.2, 0.2],
                         scope='ssd_bboxes_decode'):
    with tf.name_scope(scope):
        bboxes = []
        for i, anchors_layer in enumerate(anchors):
            bboxes.append(tf_ssd_bboxes_decode_layer(
               feat_localizations[i],
               anchors_layer,
               prior_scaling))
        return bboxes

"""
Extract classes, scores and bounding boxes from features in one layer.
Batch-compatible: inputs are supposed to have batch-type shapes.

Args:
  predictions_layer: A SSD prediction layer;
  localizations_layer: A SSD localization layer;
  select_threshold: Classification threshold for selecting a box. All boxes
    under the threshold are set to 'zero'. If None, no threshold applied.
Return:
  d_scores, d_bboxes: Dictionary of scores and bboxes Tensors of
    size Batches X N x 1 | 4. Each key corresponding to a class.
"""
def tf_ssd_bboxes_select_layer(predictions_layer, 
                               localizations_layer,
                               select_threshold=None,
                               num_classes=21,
                               ignore_class=0,
                               scope=None):

    select_threshold = 0.0 if select_threshold is None else select_threshold
    with tf.name_scope(scope, 'ssd_bboxes_select_layer',
                       [predictions_layer, localizations_layer]):
        # Reshape features: Batches x N x N_labels | 4
        p_shape = tfe.get_shape(predictions_layer)
        predictions_layer = tf.reshape(predictions_layer, tf.stack([p_shape[0], -1, p_shape[-1]]))
        l_shape = tfe.get_shape(localizations_layer)
        localizations_layer = tf.reshape(localizations_layer, tf.stack([l_shape[0], -1, l_shape[-1]]))

        d_scores = {}
        d_bboxes = {}
        for c in range(0, num_classes):
            if c != ignore_class:
                # Remove boxes under the threshold.
                scores = predictions_layer[:, :, c]
                fmask = tf.cast(tf.greater_equal(scores, select_threshold), scores.dtype)
                scores = scores * fmask
                bboxes = localizations_layer * tf.expand_dims(fmask, axis=-1)
                # Append to dictionary.
                d_scores[c] = scores
                d_bboxes[c] = bboxes

        return d_scores, d_bboxes

if __name__ == '__main__':
	
	print 'OK'






